package org.ysh.core.ocr;

import com.benjaminwan.ocrlibrary.OcrEngine;
import com.benjaminwan.ocrlibrary.OcrResult;

public class OcrUtils {


    private static final String modelsDir = "models";
    private static final String detName = "ch_PP-OCRv3_det_infer.onnx";
    private static final String clsName = "ch_ppocr_mobile_v2.0_cls_infer.onnx";
    private static final String recName = "ch_PP-OCRv3_rec_infer.onnx";
    private static final String keysName = "ppocr_keys_v1.txt";

    private static final int numThread = 8;
    private static final int padding = 50;
    private static final float boxThresh = 0.3f;
    private static final float boxScoreThresh = 0.5f;
    private static final float unClipRatio = 1.6f;
    private static final boolean doAngle = true;
    private static final boolean mostAngle = true;
    private static final int maxSideLen = 1024;

    private static final OcrEngine ocrEngine;


    static {
        ocrEngine = new OcrEngine();
        String version = ocrEngine.getVersion();
        System.out.println("version=" + version);

        ocrEngine.setNumThread(numThread);
        //------- init Logger -------
        ocrEngine.initLogger(
                false,
                false,
                false
        );
//        ocrEngine.enableResultText(imagePath);

        //------- init Models -------
        boolean initModelsRet = ocrEngine.initModels(modelsDir, detName, clsName, recName, keysName);
        if (!initModelsRet) {
            System.out.println("Error in models initialization, please check the models/keys path!");
        }else{
            //------- set param -------
            System.out.printf("padding(%d) boxScoreThresh(%f) boxThresh(%f) unClipRatio(%f) doAngle(%b) mostAngle(%b)", padding, boxScoreThresh, boxThresh, unClipRatio, doAngle, mostAngle);
            ocrEngine.setPadding(padding);
            ocrEngine.setBoxScoreThresh(boxScoreThresh);
            ocrEngine.setBoxThresh(boxThresh);
            ocrEngine.setUnClipRatio(unClipRatio);
            ocrEngine.setDoAngle(doAngle);
            ocrEngine.setMostAngle(mostAngle);
        }
    }

    /**
     * ocr图片内容
     * @param imgPath
     * @return
     */
    public static OcrResult detect(String imgPath) {

        return ocrEngine.detect(imgPath,maxSideLen);
    }


}
